//! Jaeger分布式链路追踪插件
//! 
//! 提供分布式系统的链路追踪支持，包括：
//! - OpenTelemetry集成和Jaeger导出
//! - 分布式追踪上下文传播
//! - 性能指标收集和分析
//! - 错误追踪和调试支持
//! - 服务依赖关系可视化

use anyhow::{Context, Result};
use async_trait::async_trait;
use opentelemetry::{
    global,
    trace::{TraceError, Tracer, TracerProvider},
    KeyValue,
};
use opentelemetry_sdk::{
    trace::{self, RandomIdGenerator, Sampler},
    Resource,
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tracing::{info, debug, warn, error, instrument};
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

#[cfg(test)]
mod tests;

/// Jaeger追踪配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JaegerConfig {
    /// Jaeger Agent地址
    pub agent_endpoint: String,
    /// 服务名称
    pub service_name: String,
    /// 服务版本
    pub service_version: String,
    /// 环境名称
    pub environment: String,
    /// 采样率 (0.0 - 1.0)
    pub sampling_rate: f64,
    /// 批次大小
    pub batch_size: usize,
    /// 导出超时时间（秒）
    pub export_timeout: u64,
    /// 启用控制台输出
    pub enable_console: bool,
    /// 启用JSON格式
    pub enable_json: bool,
    /// 最大属性数量
    pub max_attributes: u32,
}

impl Default for JaegerConfig {
    fn default() -> Self {
        Self {
            agent_endpoint: "http://localhost:14268/api/traces".to_string(),
            service_name: "data-gateway".to_string(),
            service_version: "0.1.0".to_string(),
            environment: "development".to_string(),
            sampling_rate: 1.0,
            batch_size: 512,
            export_timeout: 30,
            enable_console: true,
            enable_json: false,
            max_attributes: 128,
        }
    }
}

/// 追踪统计信息
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct TracingStats {
    /// 总Span数量
    pub total_spans: u64,
    /// 成功导出的Span数量
    pub exported_spans: u64,
    /// 失败的Span数量
    pub failed_spans: u64,
    /// 平均Span持续时间（毫秒）
    pub avg_span_duration_ms: f64,
    /// 最后导出时间
    pub last_export_time: Option<chrono::DateTime<chrono::Utc>>,
    /// 服务统计
    pub service_stats: HashMap<String, ServiceStats>,
}

/// 服务统计信息
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct ServiceStats {
    /// 请求数量
    pub request_count: u64,
    /// 错误数量
    pub error_count: u64,
    /// 平均响应时间（毫秒）
    pub avg_response_time_ms: f64,
    /// 最后活动时间
    pub last_activity: Option<chrono::DateTime<chrono::Utc>>,
}

/// Span信息
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SpanInfo {
    /// Span ID
    pub span_id: String,
    /// Trace ID
    pub trace_id: String,
    /// 操作名称
    pub operation_name: String,
    /// 开始时间
    pub start_time: chrono::DateTime<chrono::Utc>,
    /// 结束时间
    pub end_time: Option<chrono::DateTime<chrono::Utc>>,
    /// 持续时间（毫秒）
    pub duration_ms: Option<f64>,
    /// 标签
    pub tags: HashMap<String, String>,
    /// 状态
    pub status: SpanStatus,
}

/// Span状态
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum SpanStatus {
    /// 正在进行
    InProgress,
    /// 成功完成
    Ok,
    /// 错误
    Error,
    /// 取消
    Cancelled,
}

/// Jaeger追踪管理器
pub struct JaegerTracingManager {
    /// 配置信息
    config: JaegerConfig,
    /// 追踪器
    tracer: Box<dyn Tracer + Send + Sync>,
    /// 统计信息
    stats: Arc<RwLock<TracingStats>>,
    /// 活动Span
    active_spans: Arc<RwLock<HashMap<String, SpanInfo>>>,
}

impl JaegerTracingManager {
    /// 创建新的Jaeger追踪管理器
    pub async fn new(config: JaegerConfig) -> Result<Self> {
        info!("初始化Jaeger追踪管理器，服务: {}", config.service_name);
        
        // 初始化追踪器
        let tracer = Self::init_tracer(&config).await?;
        
        // 初始化订阅器
        Self::init_subscriber(&config)?;
        
        Ok(Self {
            config,
            tracer,
            stats: Arc::new(RwLock::new(TracingStats::default())),
            active_spans: Arc::new(RwLock::new(HashMap::new())),
        })
    }

    /// 初始化追踪器
    async fn init_tracer(config: &JaegerConfig) -> Result<Box<dyn Tracer + Send + Sync>> {
        // 创建Jaeger导出器
        let jaeger_exporter = opentelemetry_jaeger::new_agent_pipeline()
            .with_endpoint(&config.agent_endpoint)
            .with_service_name(&config.service_name)
            .with_auto_split_batch(true)
            .with_max_packet_size(65000)
            .build_exporter()
            .context("创建Jaeger导出器失败")?;

        // 创建追踪器提供者
        let tracer_provider = trace::TracerProvider::builder()
            .with_simple_exporter(jaeger_exporter)
            .with_config(
                trace::Config::default()
                    .with_sampler(Sampler::TraceIdRatioBased(config.sampling_rate))
                    .with_id_generator(RandomIdGenerator::default())
                    .with_max_events_per_span(64)
                    .with_max_attributes_per_span(config.max_attributes)
                    .with_resource(Resource::new(vec![
                        KeyValue::new("service.name", config.service_name.clone()),
                        KeyValue::new("service.version", config.service_version.clone()),
                        KeyValue::new("deployment.environment", config.environment.clone()),
                    ])),
            )
            .build();

        // 设置全局追踪器提供者
        global::set_tracer_provider(tracer_provider.clone());

        // 获取追踪器
        let tracer = tracer_provider.tracer("data-gateway");
        
        Ok(Box::new(tracer))
    }

    /// 初始化订阅器
    fn init_subscriber(config: &JaegerConfig) -> Result<()> {
        let telemetry = tracing_opentelemetry::layer().with_tracer(global::tracer("data-gateway"));
        
        let subscriber = tracing_subscriber::registry().with(telemetry);
        
        // 添加控制台输出
        if config.enable_console {
            if config.enable_json {
                let json_layer = tracing_subscriber::fmt::layer()
                    .json()
                    .with_current_span(false);
                subscriber.with(json_layer).init();
            } else {
                let fmt_layer = tracing_subscriber::fmt::layer()
                    .with_target(false)
                    .with_thread_ids(true)
                    .with_line_number(true);
                subscriber.with(fmt_layer).init();
            }
        } else {
            subscriber.init();
        }
        
        Ok(())
    }

    /// 开始新的Span
    #[instrument(skip(self))]
    pub async fn start_span(&self, operation_name: &str, tags: HashMap<String, String>) -> Result<String> {
        debug!("开始新的Span: {}", operation_name);
        
        let span_id = uuid::Uuid::new_v4().to_string();
        let trace_id = uuid::Uuid::new_v4().to_string();
        
        let span_info = SpanInfo {
            span_id: span_id.clone(),
            trace_id,
            operation_name: operation_name.to_string(),
            start_time: chrono::Utc::now(),
            end_time: None,
            duration_ms: None,
            tags,
            status: SpanStatus::InProgress,
        };
        
        // 保存活动Span
        {
            let mut active_spans = self.active_spans.write().await;
            active_spans.insert(span_id.clone(), span_info);
        }
        
        // 更新统计信息
        {
            let mut stats = self.stats.write().await;
            stats.total_spans += 1;
        }
        
        Ok(span_id)
    }

    /// 结束Span
    #[instrument(skip(self))]
    pub async fn finish_span(&self, span_id: &str, status: SpanStatus) -> Result<()> {
        debug!("结束Span: {}", span_id);
        
        let span_info = {
            let mut active_spans = self.active_spans.write().await;
            active_spans.remove(span_id)
        };
        
        if let Some(mut span_info) = span_info {
            let end_time = chrono::Utc::now();
            let duration = end_time.signed_duration_since(span_info.start_time);
            let duration_ms = duration.num_milliseconds() as f64;
            
            span_info.end_time = Some(end_time);
            span_info.duration_ms = Some(duration_ms);
            span_info.status = status.clone();
            
            // 更新统计信息
            {
                let mut stats = self.stats.write().await;
                match status {
                    SpanStatus::Ok => stats.exported_spans += 1,
                    SpanStatus::Error | SpanStatus::Cancelled => stats.failed_spans += 1,
                    _ => {}
                }
                
                // 更新平均持续时间
                if stats.exported_spans > 0 {
                    let total_duration = stats.avg_span_duration_ms * (stats.exported_spans - 1) as f64 + duration_ms;
                    stats.avg_span_duration_ms = total_duration / stats.exported_spans as f64;
                }
                
                stats.last_export_time = Some(chrono::Utc::now());
            }
            
            info!("Span完成: {} ({}ms)", span_info.operation_name, duration_ms);
        } else {
            warn!("未找到Span: {}", span_id);
        }
        
        Ok(())
    }

    /// 添加Span标签
    #[instrument(skip(self))]
    pub async fn add_span_tag(&self, span_id: &str, key: &str, value: &str) -> Result<()> {
        debug!("添加Span标签: {} = {}", key, value);
        
        let mut active_spans = self.active_spans.write().await;
        if let Some(span_info) = active_spans.get_mut(span_id) {
            span_info.tags.insert(key.to_string(), value.to_string());
        } else {
            warn!("未找到Span: {}", span_id);
        }
        
        Ok(())
    }

    /// 记录错误
    #[instrument(skip(self))]
    pub async fn record_error(&self, span_id: &str, error: &str) -> Result<()> {
        error!("记录错误: {}", error);
        
        self.add_span_tag(span_id, "error", "true").await?;
        self.add_span_tag(span_id, "error.message", error).await?;
        
        Ok(())
    }

    /// 获取追踪统计信息
    pub async fn get_stats(&self) -> TracingStats {
        self.stats.read().await.clone()
    }

    /// 获取活动Span列表
    pub async fn get_active_spans(&self) -> Vec<SpanInfo> {
        let active_spans = self.active_spans.read().await;
        active_spans.values().cloned().collect()
    }

    /// 健康检查
    pub async fn health_check(&self) -> Result<bool> {
        debug!("执行追踪系统健康检查");
        
        // 尝试创建一个测试Span
        let span_id = self.start_span("health_check", HashMap::new()).await?;
        self.finish_span(&span_id, SpanStatus::Ok).await?;
        
        info!("追踪系统健康检查通过");
        Ok(true)
    }

    /// 刷新追踪数据
    pub async fn flush(&self) -> Result<()> {
        debug!("刷新追踪数据");
        
        // 强制导出所有待处理的追踪数据
        global::shutdown_tracer_provider();
        
        info!("追踪数据刷新完成");
        Ok(())
    }

    /// 获取服务依赖关系
    pub async fn get_service_dependencies(&self) -> Result<HashMap<String, Vec<String>>> {
        debug!("获取服务依赖关系");
        
        // 这里应该从Jaeger查询API获取实际的依赖关系
        // 为了演示，我们返回一个示例依赖图
        let mut dependencies = HashMap::new();
        dependencies.insert("data-gateway".to_string(), vec![
            "mysql".to_string(),
            "redis".to_string(),
            "elasticsearch".to_string(),
        ]);
        dependencies.insert("mysql".to_string(), vec![]);
        dependencies.insert("redis".to_string(), vec![]);
        dependencies.insert("elasticsearch".to_string(), vec![]);
        
        Ok(dependencies)
    }

    /// 查询追踪数据
    pub async fn query_traces(&self, service: Option<&str>, operation: Option<&str>, limit: usize) -> Result<Vec<SpanInfo>> {
        debug!("查询追踪数据: service={:?}, operation={:?}, limit={}", service, operation, limit);
        
        // 这里应该从Jaeger查询API获取实际的追踪数据
        // 为了演示，我们返回当前活动的Span
        let active_spans = self.get_active_spans().await;
        
        let filtered_spans: Vec<SpanInfo> = active_spans
            .into_iter()
            .filter(|span| {
                if let Some(svc) = service {
                    if !span.tags.get("service.name").map_or(false, |s| s == svc) {
                        return false;
                    }
                }
                if let Some(op) = operation {
                    if span.operation_name != op {
                        return false;
                    }
                }
                true
            })
            .take(limit)
            .collect();
        
        Ok(filtered_spans)
    }

    /// 获取性能指标
    pub async fn get_performance_metrics(&self) -> Result<HashMap<String, f64>> {
        debug!("获取性能指标");
        
        let stats = self.get_stats().await;
        let mut metrics = HashMap::new();
        
        metrics.insert("total_spans".to_string(), stats.total_spans as f64);
        metrics.insert("exported_spans".to_string(), stats.exported_spans as f64);
        metrics.insert("failed_spans".to_string(), stats.failed_spans as f64);
        metrics.insert("avg_span_duration_ms".to_string(), stats.avg_span_duration_ms);
        
        if stats.total_spans > 0 {
            let success_rate = stats.exported_spans as f64 / stats.total_spans as f64 * 100.0;
            metrics.insert("success_rate_percent".to_string(), success_rate);
        }
        
        Ok(metrics)
    }
}

/// Jaeger追踪服务特征
#[async_trait]
pub trait JaegerTracingService {
    /// 开始Span
    async fn start_span(&self, operation_name: &str, tags: HashMap<String, String>) -> Result<String>;
    
    /// 结束Span
    async fn finish_span(&self, span_id: &str, status: SpanStatus) -> Result<()>;
    
    /// 添加标签
    async fn add_span_tag(&self, span_id: &str, key: &str, value: &str) -> Result<()>;
    
    /// 记录错误
    async fn record_error(&self, span_id: &str, error: &str) -> Result<()>;
    
    /// 获取统计信息
    async fn get_stats(&self) -> TracingStats;
    
    /// 健康检查
    async fn health_check(&self) -> Result<bool>;
}

#[async_trait]
impl JaegerTracingService for JaegerTracingManager {
    async fn start_span(&self, operation_name: &str, tags: HashMap<String, String>) -> Result<String> {
        self.start_span(operation_name, tags).await
    }
    
    async fn finish_span(&self, span_id: &str, status: SpanStatus) -> Result<()> {
        self.finish_span(span_id, status).await
    }
    
    async fn add_span_tag(&self, span_id: &str, key: &str, value: &str) -> Result<()> {
        self.add_span_tag(span_id, key, value).await
    }
    
    async fn record_error(&self, span_id: &str, error: &str) -> Result<()> {
        self.record_error(span_id, error).await
    }
    
    async fn get_stats(&self) -> TracingStats {
        self.get_stats().await
    }
    
    async fn health_check(&self) -> Result<bool> {
        self.health_check().await
    }
}
